{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Habitat 2.0 Gym API\n",
    "This tutorial covers how to use Habitat 2.0 environments as standard gym environments.\n",
    "See [here for Habitat 2.0 installation instructions and more tutorials.](https://aihabitat.org/docs/habitat2/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import git\n",
    "\n",
    "if \"COLAB_GPU\" in os.environ:\n",
    "    print(\"Setting Habitat base path\")\n",
    "    %env HABLAB_BASE_CFG_PATH=/content/habitat-lab\n",
    "    import importlib\n",
    "\n",
    "    import PIL\n",
    "\n",
    "    importlib.reload(PIL.TiffTags)  # type: ignore[attr-defined]\n",
    "\n",
    "import imageio\n",
    "\n",
    "# Video rendering utility.\n",
    "from habitat_sim.utils import viz_utils as vut\n",
    "\n",
    "# Quiet the Habitat simulator logging\n",
    "os.environ[\"MAGNUM_LOG\"] = \"quiet\"\n",
    "os.environ[\"HABITAT_SIM_LOG\"] = \"quiet\"\n",
    "\n",
    "repo = git.Repo(\".\", search_parent_directories=True)\n",
    "dir_path = repo.working_tree_dir\n",
    "output_path = os.path.join(\n",
    "    dir_path, \"examples/tutorials/habitat_lab_visualization/\"\n",
    ")\n",
    "os.makedirs(output_path, exist_ok=True)\n",
    "os.chdir(dir_path)\n",
    "# If the import block below fails due to an error like \"'PIL.TiffTags' has no attribute\n",
    "# 'IFD'\", then restart the Colab runtime instance and rerun this cell and the previous cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The ONLY two lines you need to add to start importing Habitat 2.0 Gym environments.\n",
    "import gym\n",
    "\n",
    "# flake8: noqa\n",
    "import habitat.gym"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simple Example\n",
    "This example sets up the Pick task in render mode which includes a high resolution camera in the scene for visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(\"HabitatRenderPick-v0\")\n",
    "\n",
    "video_file_path = os.path.join(output_path, \"example_interact.mp4\")\n",
    "video_writer = imageio.get_writer(video_file_path, fps=30)\n",
    "\n",
    "done = False\n",
    "env.reset()\n",
    "while not done:\n",
    "    obs, reward, done, info = env.step(env.action_space.sample())\n",
    "    video_writer.append_data(env.render(\"rgb_array\"))\n",
    "\n",
    "video_writer.close()\n",
    "if vut.is_notebook():\n",
    "    vut.display_video(video_file_path)\n",
    "\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Environment Options\n",
    "To create the environment in performance mode remove `Render` from the environment ID string. The environment ID follows the format: `Habitat[Render?][Task Name]-v0`. All the supported environment IDs are listed below. The `Render` option can always be added to include the higher resolution 3rd POV camera for visualization.\n",
    "\n",
    "* Skills:\n",
    "    * `HabitatPick-v0`\n",
    "    * `HabitatPlace-v0`\n",
    "    * `HabitatCloseCab-v0`\n",
    "    * `HabitatCloseFridge-v0`\n",
    "    * `HabitatOpenCab-v0`\n",
    "    * `HabitatOpenFridge-v0`\n",
    "    * `HabitatNavToObj-v0`\n",
    "    * `HabitatReachState-v0`\n",
    "* Home Assistant Benchmark (HAB) tasks:\n",
    "    * `HabitatTidyHouse-v0`\n",
    "    * `HabitatPrepareGroceries-v0`\n",
    "    * `HabitatSetTable-v0`\n",
    "\n",
    "The Gym environments are automatically registered from the RL training configurations under [\"habitat-lab/habitat/config/benchmark/rearrange\"](https://github.com/facebookresearch/habitat-lab/tree/main/habitat-lab/habitat/config/benchmark/rearrange). The observation keys in `habitat.gym.obs_keys` are what is returned in the observation space.\n",
    "\n",
    "An example of these different observation spaces is demonstrated below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dictionary observation space\n",
    "env = gym.make(\"HabitatPick-v0\")\n",
    "print(\n",
    "    \"Pick observation space\",\n",
    "    {k: v.shape for k, v in env.observation_space.spaces.items()},\n",
    ")\n",
    "env.close()\n",
    "\n",
    "# Array observation space\n",
    "env = gym.make(\"HabitatReachState-v0\")\n",
    "print(\"Reach observation space\", env.observation_space)\n",
    "env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Environment Configuration\n",
    "\n",
    "You can also modify the config specified in the YAML file through `gym.make` by passing the `override_options` argument. Here is an example of changing the gripper type to use the suction grasp in the Pick Task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make(\n",
    "    \"HabitatPick-v0\",\n",
    "    override_options=[\n",
    "        \"habitat.task.actions.arm_action.grip_controller=SuctionGraspAction\",\n",
    "    ],\n",
    ")\n",
    "print(\"Action space with suction grip\", env.action_space)\n",
    "env.close()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Habitat 2.0 Gym Tutorial",
   "provenance": []
  },
  "jupytext": {
   "cell_metadata_filter": "-all",
   "formats": "nb_python//py:percent,notebooks//ipynb",
   "notebook_metadata_filter": "all"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
